from pymilvus import connections, FieldSchema, CollectionSchema, DataType, Collection
import numpy as np

# 连接到 Milvus 集群（目前不支持window，需要部署linux或docker）
connections.connect(
    alias="default",
    host='localhost',
    port='19530'
)

# 定义集合的字段
fields = [
    FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
    FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=128)
]

# 定义集合的模式
schema = CollectionSchema(fields=fields, description="Test collection")

# 创建集合
collection_name = "test_collection"
collection = Collection(name=collection_name, schema=schema)

# 生成一些示例向量
vectors = np.random.rand(10, 128).astype(np.float32)

# 插入向量到集合中
data = [vectors]
collection.insert(data)

# 构建索引
index = {
    "index_type": "IVF_FLAT",
    "metric_type": "L2",
    "params": {"nlist": 128},
}
collection.create_index(field_name="embedding", index_params=index)

# 加载集合
collection.load()

# 进行向量搜索
search_params = {
    "metric_type": "L2",
    "params": {"nprobe": 10},
}
query_vector = np.random.rand(1, 128).astype(np.float32)
results = collection.search(
    data=query_vector,
    anns_field="embedding",
    param=search_params,
    limit=10,
    output_fields=["id"]
)

# 打印搜索结果
for hit in results[0]:
    print(f"ID: {hit.id}, Distance: {hit.distance}")

# 断开连接
connections.disconnect("default")
